package main

import (
	"context"
	"fmt"
	"log"
	"time"

	clientv3 "go.etcd.io/etcd/client/v3"
	"go.etcd.io/etcd/client/v3/concurrency"
)

type EtcdLock struct {
	session *concurrency.Session
	mutex   *concurrency.Mutex
}

func NewEtcdLock(endpoints []string, lockKey string) (*EtcdLock, error) {
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   endpoints,
		DialTimeout: 5 * time.Second,
	})
	if err != nil {
		return nil, err
	}

	session, err := concurrency.NewSession(cli, concurrency.WithTTL(5))
	if err != nil {
		cli.Close()
		return nil, err
	}

	mutex := concurrency.NewMutex(session, lockKey)

	return &EtcdLock{
		session: session,
		mutex:   mutex,
	}, nil
}

func (l *EtcdLock) Lock(ctx context.Context) error {
	return l.mutex.Lock(ctx)
}

func (l *EtcdLock) Unlock(ctx context.Context) error {
	return l.mutex.Unlock(ctx)
}

func (l *EtcdLock) Close() {
	l.session.Close()
}

func main() {
	endpoints := []string{"localhost:2379"}
	lockKey := "/my-distributed-lock"

	lock, err := NewEtcdLock(endpoints, lockKey)
	if err != nil {
		log.Fatalf("failed to create lock: %v", err)
	}
	defer lock.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()

	err = lock.Lock(ctx)
	if err != nil {
		log.Fatalf("failed to acquire lock: %v", err)
	}
	fmt.Println("Lock acquired")

	// Simulate some work with the lock
	time.Sleep(10 * time.Second)

	err = lock.Unlock(ctx)
	if err != nil {
		log.Fatalf("failed to release lock: %v", err)
	}
	fmt.Println("Lock released")
}
